import os
import torch
import yaml
from model import PoetryModel
from utils import load_dataset, generate_poem


class Validator:
	def __init__(self, config, model_path=None):
		"""初始化验证器，支持从指定路径加载模型"""
		self.config = config
		self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

		# 加载数据集以获取字典
		self.word2ix, self.ix2word, _ = load_dataset(config['data']['path'])
		self.vocab_size = len(self.word2ix)

		# 创建模型
		self.model = PoetryModel(
			vocab_size=self.vocab_size,
			d_model=config['model']['d_model'],
			nhead=config['model']['nhead'],
			num_layers=config['model']['num_layers'],
			dim_feedforward=config['model']['dim_feedforward'],
			max_seq_len=config['model']['max_seq_len']
		).to(self.device)

		# 加载模型权重
		if model_path:
			# 使用指定路径的模型
			self.load_model(model_path)
		else:
			# 使用最佳模型（从配置的保存目录加载）
			self.load_best_model()

		# 设置生成参数
		self.max_gen_len = config['validation']['max_gen_len']
		self.temperature = config['validation']['temperature']
		self.repetition_penalty = config['validation']['repetition_penalty']

	def load_best_model(self):
		"""加载最优模型（从保存目录中选择）"""
		model_dir = self.config['model']['save_dir']
		history_file = os.path.join(model_dir, 'training_history.yaml')

		if not os.path.exists(history_file):
			raise FileNotFoundError("未找到训练历史文件，请先训练模型或指定模型路径")

		with open(history_file, 'r', encoding='utf-8') as f:
			history = yaml.safe_load(f)

		if not history:
			raise ValueError("训练历史文件为空")

		# 找到损失最小的epoch
		best_epoch = min(history, key=lambda x: x['train_loss'])['epoch']
		model_path = os.path.join(model_dir, f'model_epoch_{best_epoch}.pth')

		# 加载模型权重
		self.load_model(model_path)
		print(f"已加载最佳模型 (epoch {best_epoch})")

	def load_model(self, path):
		"""加载指定路径的模型"""
		if not os.path.exists(path):
			raise FileNotFoundError(f"模型文件不存在: {path}")

		checkpoint = torch.load(path, map_location=self.device)
		self.model.load_state_dict(checkpoint['model_state_dict'])
		self.model.eval()
		print(f"已加载模型: {path}")

	def validate(self):
		"""执行验证，生成古诗"""
		print("===== 古诗生成系统 =====")

		while True:
			prompt = input("请输入开头（直接回车则随机生成，输入q退出）：")
			if prompt.lower() == 'q':
				break

			poem = generate_poem(
				self.model,
				prompt,
				self.ix2word,
				self.word2ix,
				self.max_gen_len,
				self.temperature,
				self.repetition_penalty,
				self.device
			)

			print("\n生成的古诗：")
			print(poem)
			print("=" * 30)